import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models

class VGGNet(nn.Module):
    """
    VGG网络实现，基于torchvision的VGG16模型
    """
    def __init__(self, pretrained=True):
        super(VGGNet, self).__init__()
        # 加载预训练的VGG16模型
        # 使用新的weights参数代替pretrained
        if pretrained:
            vgg16 = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1)
        else:
            vgg16 = models.vgg16(weights=None)
        
        # 提取特征部分 (卷积层)
        self.features = vgg16.features

        # avgpool
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        
        # 提取分类器部分 (全连接层)，但不包括最后一层
        self.classifier = nn.Sequential(*list(vgg16.classifier.children())[:-1])
        
        # 添加我们自己的输出层，对应于Caffe中的fc8_m
        self.fc8_m = nn.Linear(4096, 1)
        
    def forward(self, x):
        # 特征提取
        x = self.features(x)

        # avgpool
        x = self.avgpool(x)
        
        # 展平特征图
        x = x.view(x.size(0), -1)
        
        # 全连接层
        x = self.classifier(x)
        
        # 最终输出层
        x = self.fc8_m(x)
        
        return x
    
    def get_parameter_groups(self):
        """
        获取参数组，模拟Caffe中的lr_mult和decay_mult
        
        在Caffe中：
        - 权重的lr_mult=1, decay_mult=1
        - 偏置的lr_mult=2, decay_mult=0
        - fc8_m层的lr_mult=10(权重)/20(偏置), decay_mult=1/0
        
        返回:
            参数组列表，可直接用于优化器
        """
        # 基础层参数 (除了fc8_m)
        base_params = []
        # fc8_m层参数 (最后一层)
        fc8_params = []
        
        # 遍历所有命名参数
        for name, param in self.named_parameters():
            # 根据参数名称判断是否为fc8_m层
            if 'fc8_m' in name:
                if 'weight' in name:
                    # fc8_m权重 - lr_mult=10, decay_mult=1
                    fc8_params.append({'params': param, 'lr': 0.001, 'weight_decay': 0.0005})
                else:
                    # fc8_m偏置 - lr_mult=20, decay_mult=0
                    fc8_params.append({'params': param, 'lr': 0.002, 'weight_decay': 0})
            else:
                if 'weight' in name:
                    # 其他层权重 - lr_mult=1, decay_mult=1
                    base_params.append({'params': param, 'lr': 0.0001, 'weight_decay': 0.0005})
                else:
                    # 其他层偏置 - lr_mult=2, decay_mult=0
                    base_params.append({'params': param, 'lr': 0.0002, 'weight_decay': 0})
        
        # 返回两组参数
        return base_params + fc8_params
    